package net.tp.algo.heap;

import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.NoSuchElementException;

/**
 * 
 * <p>Binary Min-Heap. Binary min-heap is a complete binary tree where the parent element is smaller than or equal to its two children.
 * The root of the tree is the minimum element. To construct a max-heap, pass in a reverse comparator.</p>
 * 
 * <p>Making binary heap from an array of length n takes $O(n)$ time. The strategy is to go through each element in the upwards
 * order and sift it down the tree. Assume the tree is a full binary tree, the number of leaves is ${n+1}/2$.</p>
 * 
 * <p>Let $h$ be the height of an element from the bottom. The cost of sifting down an element is $O(h)$. The operations for all
 * the elements at the leaf level is ${n+1}/2 O(1) = n/2 O(1)$. The next level up is ${n+1}/2^2 O(2) = n/2^2 O(2)$. So the total running time is:
 * $O(n (1/2^1 + 2/2^2 + 3/2^3 + ... + {log(n)}/2^{log(n)} ))$. Note: $n/2^n &lt; 1/n^2$ for all sufficiently large $n$.
 * So for some constant $C$:</p>
 * $$1/2^1 + 2/2^2 + &hellip; {log(n)}/2^{log(n)} &lt; 1/2^1 + 2/2^2 + &hellip; n/2^n &lt; C + 1/1^2 + 1/2^2 + 1/3^2 + &hellip; + 1/n^2 &rarr; C + &pi;^2/6 $$
 * <p>See <a href="http://en.wikipedia.org/wiki/Basel_problem">Basel problem</a>. So the running time of making heap is $O(n)$.</p>
 * 
 * @author Trung Phan
 *
 * @param <E> element type
 * @time construct, add, remove, update: $O(n)$. Read: $O(n)$.
 * @space $O(n)$.
 */
public class BinaryHeap<E> {
	
	private E[] elements;
	private int n;
	private final Comparator<? super E> comparator;
	
	private transient Map<E, Integer> elementPositions;
	
	/**
	 * Create min-heap.
	 */
	public BinaryHeap() {
		this((Collection<E>)null, null);
	}
	
	/**
	 * Create min-heap. Reverse comparator for max-heap.
	 * @param comparator comparator for elements.
	 */
	public BinaryHeap(Comparator<? super E> comparator) {
		this(null, comparator);
	}
	
	/**
	 * Create min-heap with predefined elements.
	 * 
     * @param elements initial elements
	 * @see BinaryHeap#BinaryHeap(Collection, Comparator)
	 */
	public BinaryHeap(Collection<E> elements) {
		this(elements, null);
	}
	
	/**
	 * Create min-heap with predefined elements.
	 * 
     * @param elements initial elements
	 * @see BinaryHeap#BinaryHeap(Collection, Comparator)
	 */
	@SafeVarargs
    @SuppressWarnings("varargs")
	public BinaryHeap(E ... elements) {
		this(Arrays.asList(elements), null);
	}
	
	/**
	 * Create min-heap with predefined elements. Reverse comparator for max-heap.
	 * 
     * @param comparator element comparator
     * @param elements initial elements
	 * @see BinaryHeap#BinaryHeap(Collection, Comparator)
	 */
	@SafeVarargs
    @SuppressWarnings("varargs")
	public BinaryHeap(Comparator<? super E> comparator, E ... elements) {
		this(Arrays.asList(elements), comparator);
	}
	
	/**
	 * Create min-heap with predefined elements. Reverse comparator for max-heap.
	 * 
	 * @time $O(n log(n))$
	 * @space $O(n)$
	 * @param elements Ignored if null.
	 * @param comparator comparator for elements. If null, elements must be comparable.
	 */
	@SuppressWarnings("unchecked")
	public BinaryHeap(Collection<E> elements, Comparator<? super E> comparator) {
		this.comparator = comparator;
		
		if (elements == null) {
			this.elements = (E[])new Object[16];
		}
		else {
			this.n = elements.size();
			int len = Math.max(16, Integer.highestOneBit(n) * 2);
			
			this.elements = (E[])new Object[len];
			int i = 0;
			for (E element : elements) {
				this.elements[i++] = element;
			}
			
			heapify();
		}
	}
	
	/**
	 * Return true if the heap is empty.
	 * @return true if the heap is empty.
	 * @time $O(1)$
	 */
	public boolean empty() {
		return n == 0;
	}
	
	/**
	 * Size of the heap.
	 * @return size of the heap.
	 * @time $O(1)$
	 */
	public int size() {
		return n;
	}

	/**
	 * Return minimum element for min-heap, maximum element for max-heap. The element is not removed about of the heap.
	 * 
	 * @time $O(1)$
	 * @return minimum element for min-heap, maximum element for max-heap.
	 * @throws NoSuchElementException if heap is empty.
	 */
	public E head() {
		if (n == 0) throw new NoSuchElementException();
		return elements[0];
	}
	
	/**
	 * Remove the minimum element (if min-heap).
	 * 
	 * @time $O(n log(n))$
	 * @return minimum element for min-heap, maximum element for max-heap.
	 * @throws NoSuchElementException if heap is empty.
	 */
	public E removeHead() {
		if (n == 0) throw new NoSuchElementException();
		
		E result = elements[0];
		swap(0, n-1);
		elements[n-1] = null;
		n--;
        if (elementPositions != null) elementPositions.remove(result);
		
		if (n > 0) siftDown(0, n);
		
		shrinkArrayIfPossible();
		
		return result;
	}

	/**
	 * Add an element into the heap.
	 * 
	 * @time $O(n log(n))$
	 * @param element element to add. Non null.
	 * @throws NullPointerException if element is null.
	 */
	public void add(E element) {
		if (element == null) throw new NullPointerException();
		
		expandArrayIfFull();
		
		n++;
		elements[n-1] = element;
        if (elementPositions != null) elementPositions.put(element, n-1);
		siftUp(n-1);
	}

    public boolean contains(E element) {
        initElementPositionsIfNecessary();

        Integer pos = elementPositions.get(element);

        return pos != null;
    }

	/**
	 * Remove an element from the heap. The first call may take $O(n)$ time to initialize internal structure
	 * to keep track of element's position. The subsequent call will take $O(n log(n))$ time.
	 * 
	 * @time $O(n log(n))$
	 * @param element element to remove
	 * @throws NoSuchElementException if the heap does not contain the element.
	 */
	public void remove(E element) {
		int pos = getElementPos(element);
		
		swap(pos, n-1);
        if (elementPositions != null) elementPositions.remove(elements[n-1]);
		elements[n-1] = null;
		n--;
		
		if (n > 0 && pos < n) {
			updatePriority(pos);
		}
		
		shrinkArrayIfPossible();
	}

	/**
	 * Update heap if the element's priority is changed. The first call may take $O(n)$ time to initialize internal structure
	 * to keep track of element's position. The subsequent call will take $O(n log(n))$ time.
	 * 
	 * @time $O(n log(n))$
	 * @param element element to update priority
	 */
	public void updatePriority(E element) {
		int pos = getElementPos(element);
		updatePriority(pos);
	}

	private void heapify() {
		if (n > 0) {
			for (int i = n / 2; i >= 0; i--) {
				siftDown(i, n);
			}
		}
	}
	
	private void siftDown(int i, int len) {
		assert i >= 0 && i < n;
		
		while (true) {
			int child = i * 2 + 1;
			
			if (child >= len) {
				break;
			}
			
			if (child < len - 1 && less(elements[child+1], elements[child])) {
				child++;
			}
			if (less(elements[child], elements[i])) {
				swap(i, child);
				i = child;
			}
			else {
				break;
			}
		}
	}
	
	private void siftUp(int i) {
		assert i >= 0 && i < n;
		
		while (i > 0) { // ignore root
			int parent = (i - 1) / 2; // this formula is wrong if i == 0;
			if (less(elements[i], elements[parent])) {
				swap(i, parent);
				i = parent;
			}
			else {
				break;
			}
		}
	}
	
	private void updatePriority(int pos) {
		if (pos == 0 || less(elements[(pos-1)/2], elements[pos] )) {
			siftDown(pos, n);
		}
		else {
			siftUp(pos);
		}
	}

	private void swap(int i, int j) {
		if (i != j) {
			E v = elements[i];
			elements[i] = elements[j];
			elements[j] = v;
			
			if (elementPositions != null) {
				elementPositions.put(elements[i], i);
				elementPositions.put(elements[j], j);
			}
		}
	}
	
	@SuppressWarnings("unchecked")
	private boolean less(E element1, E element2) {
		int compare = comparator != null ? comparator.compare(element1, element2):
			((Comparable<E>)element1).compareTo(element2);
		
		return compare < 0;
	}
	
	private int getElementPos(E element) {
		initElementPositionsIfNecessary();
		
		Integer pos = elementPositions.get(element);
		
		if (pos == null) {
			throw new NoSuchElementException();
		}
		return pos;
	}

	private void initElementPositionsIfNecessary() {
		if (elementPositions == null) {
			elementPositions = new HashMap<>();
			for (int i = 0; i < n; i++) {
				elementPositions.put(elements[i], i);
			}
		}
	}
	
	private void expandArrayIfFull() {
		if (n == elements.length) {
			elements = Arrays.copyOf(elements, n * 2);
		}
	}
	
	private void shrinkArrayIfPossible() {
		if (elements.length > 16 && n < elements.length / 4) {
			int newLen = Math.max(16, Integer.highestOneBit(n) * 2 );
			elements = Arrays.copyOf(elements, newLen);
		}
	}
	
	/**
	 * Return true if internal heap is valid.
	 * @return true if internal heap is valid.
	 */
	boolean valid() {
		for (int i = n - 1; i >= 1; i--) {
			int p = (i - 1) / 2;
			if (less(elements[i], elements[p])) {
				return false;
			}
		}
		return true;
	}
	
}
